//
// Created by hjie on 23-7-30.
//
#include <cstring>
#include <string>
#include <iostream>
#include <cstdio>
#include <vector>
#include <sstream>
#include "base/dump.h"
extern "C"
{
#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libavutil/avutil.h>
#include <libavutil/timestamp.h>
#include <libavutil/mathematics.h>
#include <libavutil/timestamp.h>
#include <libavutil/time.h>
#include <libavutil/opt.h>
#include <libswresample/swresample.h>
}

struct VideoAudioParameter
{
    int copy_video;
    int copy_audio;
    std::string video_encoder_codec_name;
    std::string audio_encoder_codec_name;
};

struct InputFile
{
    std::string filename;
    AVFormatContext * format_ctx;
    int video_index{0};
    int audio_index{0};
    AVCodecContext * video_decoder_ctx;
    AVCodecContext * audio_decoder_ctx;
    AVPacket * m_read_packet;
    AVFrame  * m_decoder_frame;
};

struct OutputFile
{
    std::string output_filename;
    AVFormatContext * format_ctx;
    AVCodecContext * video_encoder_ctx;
    AVCodecContext * audio_encoder_ctx;
    int video_encoder_index{0};
    int audio_encoder_index{0};
    int64_t m_next_pts{0};
    FileDumpData * m_dump_handle;

    AVPacket * m_encoder_packet;
    AVFrame  * m_frame;
};

void DumpLog(const std::string & func, const std::string & content, int error_code = 0)
{
    printf("func : %s, content : %s code : %d.\n", func.c_str(), content.c_str(), error_code);
}

static void DumpPacket(InputFile & input, OutputFile & output, AVPacket * packet)
{
    std::stringstream ss;
    std::string media_type;
    if (packet->stream_index == input.audio_index)
    {
        media_type = "audio";
        packet->pts = av_rescale_q(packet->pts, input.format_ctx->streams[packet->stream_index]->time_base, input.format_ctx->streams[input.video_index]->time_base);
        packet->dts = packet->pts;
    }
    else if (packet->stream_index == input.video_index)
    {
        media_type = "video";
    }
    ss << "media type : " << media_type << "\tpacket pts : " << packet->pts << "\tdts : " << packet->dts << "\tdata size : " << packet->size;
    output.m_dump_handle->WriteData(ss.str());
}

static void DumpFrame(FileDumpData * handle, AVFrame * frame)
{
    std::stringstream ss;
    ss << "frame pts : " << frame->pts << "\tpkt_dts : " << frame->pkt_dts;
    handle->WriteData(ss.str());
}
// open_input file
// find beast stream information
//
int OpenInputFile(InputFile & input)
{
    input.format_ctx = nullptr;
    input.format_ctx = avformat_alloc_context();
    int result = avformat_open_input(&input.format_ctx, input.filename.c_str(), nullptr, nullptr);
    if (result < 0)
    {
        DumpLog(__func__, "call avformat_open_input error", result);
        return -1;
    }
    DumpLog(__func__ , "stream number ", (int)input.format_ctx->nb_streams);
    avformat_find_stream_info(input.format_ctx, nullptr);
    // todo will add to note
    av_dump_format(input.format_ctx, 0, input.filename.c_str(), 0);

    for (auto index = 0; index < input.format_ctx->nb_streams; index++)
    {
        AVStream * stream = input.format_ctx->streams[index];
        const AVCodec * codec = avcodec_find_decoder(stream->codecpar->codec_id);
        if (stream->codecpar->codec_type == AVMEDIA_TYPE_VIDEO)
        {
            input.video_index = index;
            input.video_decoder_ctx = avcodec_alloc_context3(codec);
            avcodec_parameters_to_context(input.video_decoder_ctx, stream->codecpar);
            input.video_decoder_ctx->pkt_timebase = stream->time_base;
            result = avcodec_open2(input.video_decoder_ctx, codec, nullptr);
            if (result < 0)
            {
                DumpLog(__func__ , "call avcodec_open2 fail ", result);
                result = -1;
                break;
            }
        }
        else if (stream->codecpar->codec_type == AVMEDIA_TYPE_AUDIO)
        {
            input.audio_index = index;
            input.audio_decoder_ctx = avcodec_alloc_context3(codec);
            avcodec_parameters_to_context(input.audio_decoder_ctx, stream->codecpar);

            input.audio_decoder_ctx->pkt_timebase = stream->time_base;
            result = avcodec_open2(input.audio_decoder_ctx, codec, nullptr);
            if (result < 0)
            {
                DumpLog(__func__ , "call avcodec_open2 fail ", result);
                result = -1;
                break;
            }
        }
    }
    input.m_read_packet = av_packet_alloc();
    input.m_decoder_frame = av_frame_alloc();
    return 0;
}

void FreeInput(InputFile & input)
{
    if (input.format_ctx != nullptr)
    {
        avformat_free_context(input.format_ctx);
        input.format_ctx = nullptr;
    }
    if (input.video_decoder_ctx != nullptr)
    {
        avcodec_free_context(&input.video_decoder_ctx);
        input.video_decoder_ctx = nullptr;
    }
    if (input.audio_decoder_ctx != nullptr)
    {
        avcodec_free_context(&input.audio_decoder_ctx);
        input.audio_decoder_ctx = nullptr;
    }
    if (input.m_read_packet)
    {
        av_packet_unref(input.m_read_packet);
        av_packet_free(&input.m_read_packet);
    }
    if (input.m_decoder_frame)
    {
        av_frame_unref(input.m_decoder_frame);
        av_frame_free(&input.m_decoder_frame);
    }
    input.video_index = -1;
    input.audio_index = -1;
    input.filename.clear();
}

int OpenOutputFile(InputFile & input, OutputFile & output, VideoAudioParameter & param_set)
{
    output.format_ctx = avformat_alloc_context();
    int result = avformat_alloc_output_context2(&output.format_ctx, nullptr, "mp4", output.output_filename.c_str());
    if (result < 0)
    {
        DumpLog(__func__ , "call avformat_alloc_output_context2 fail", result);
        return -1;
    }
    output.format_ctx->duration = input.format_ctx->duration;
    output.format_ctx->start_time = input.format_ctx->start_time;

    for (auto index = 0; index < input.format_ctx->nb_streams; index++)
    {
        AVStream * stream = input.format_ctx->streams[index];
        if (stream && stream->codecpar->codec_type == AVMEDIA_TYPE_VIDEO)
        {
            AVStream * video_stream = avformat_new_stream(output.format_ctx, nullptr);
            if (video_stream)
            {
                const AVCodec * codec = avcodec_find_encoder(AV_CODEC_ID_H264);
                if (!codec)
                {
                    return -1;
                }
                output.video_encoder_ctx = avcodec_alloc_context3(codec);
                // copy param to new video encoder ctx
                avcodec_parameters_to_context(output.video_encoder_ctx, stream->codecpar);
                //output.video_encoder_ctx->width = 640;
                //output.video_encoder_ctx->height = 480;
                //output.video_encoder_ctx->sample_aspect_ratio = AVRational{4,3};
                /* video time_base can be set to whatever is handy and supported by encoder */
                //output.video_encoder_ctx->has_b_frames = 1;
                //output.video_encoder_ctx->max_b_frames = 2;
                output.video_encoder_ctx->time_base    = stream->time_base;
                output.video_encoder_ctx->bit_rate     = input.video_decoder_ctx->bit_rate / 2;
                //output.video_encoder_ctx->time_base = (AVRational){1, 2970}; // 帧与帧之间的间隔是time_base
                output.video_encoder_ctx->framerate = stream->avg_frame_rate;
                if (output.format_ctx->oformat->flags & AVFMT_GLOBALHEADER)
                {
                    output.video_encoder_ctx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
                }
                // todo may be could custom set param
                result = avcodec_open2(output.video_encoder_ctx, codec, nullptr);
                if (result < 0)
                {
                    DumpLog(__func__ , "call avcodec_open2", result);
                    result = -1;
                    break;
                }
                avcodec_parameters_from_context(video_stream->codecpar, output.video_encoder_ctx);
                video_stream->time_base = output.video_encoder_ctx->time_base;
                output.video_encoder_index = video_stream->index;
                video_stream->avg_frame_rate = stream->avg_frame_rate;
                video_stream->r_frame_rate = stream->r_frame_rate;

                output.video_encoder_index = video_stream->index;
                avcodec_flush_buffers(output.video_encoder_ctx);
            }
        }
        else if (stream && stream->codecpar->codec_type == AVMEDIA_TYPE_AUDIO)
        {
            AVStream * audio_stream = avformat_new_stream(output.format_ctx, nullptr);
            if (audio_stream)
            {
                const AVCodec * codec = avcodec_find_encoder(AV_CODEC_ID_AAC);
                output.audio_encoder_ctx = avcodec_alloc_context3(codec);
                // copy param to new video encoder ctx
                output.audio_encoder_ctx->time_base = stream->time_base;
                avcodec_parameters_to_context(output.audio_encoder_ctx, stream->codecpar);
                result = avcodec_open2(output.audio_encoder_ctx, codec, nullptr);
                if (result < 0)
                {
                    DumpLog(__func__ , "call avcodec_open2", result);
                    result = -1;
                    break;
                }
                avcodec_parameters_from_context(audio_stream->codecpar, output.audio_encoder_ctx);
                output.audio_encoder_index = audio_stream->index;
                avcodec_flush_buffers(output.audio_encoder_ctx);
            }
        }
    }
    if (!(output.format_ctx->oformat->flags & AVFMT_NOFILE))
    {
        result = avio_open(&output.format_ctx->pb, output.output_filename.c_str(), AVIO_FLAG_WRITE);
        if (result < 0)
        {
            DumpLog(__func__ , "call avio_open fail", result);
            return -1;
        }
    }
    output.m_dump_handle = new FileDumpData();
    output.m_dump_handle->InitDump();
    output.m_encoder_packet = av_packet_alloc();
    output.m_frame          = av_frame_alloc();
    return 0;
}

void FreeOutput(OutputFile & output)
{
    if (output.video_encoder_ctx != nullptr)
    {
        avcodec_free_context(&output.video_encoder_ctx);
        output.video_encoder_ctx = nullptr;
    }
    if (output.audio_encoder_ctx != nullptr)
    {
        avcodec_free_context(&output.audio_encoder_ctx);
        output.audio_encoder_ctx = nullptr;
    }
    if (output.format_ctx != nullptr)
    {
        if (!(output.format_ctx->oformat->flags & AVFMT_NOFILE))
        {
            avio_close(output.format_ctx->pb);
        }
        avformat_free_context(output.format_ctx);
        output.format_ctx = nullptr;
    }
    if (output.m_dump_handle)
    {
        delete output.m_dump_handle;
    }
    if (output.m_encoder_packet)
    {
        av_packet_unref(output.m_encoder_packet);
        av_packet_free(&output.m_encoder_packet);
    }
    if (output.m_frame)
    {
        av_frame_unref(output.m_frame);
        av_frame_free(&output.m_frame);
    }
    output.output_filename.clear();
}

int ReadInputFrame(InputFile & input, AVPacket * packet)
{
    // todo push packet to the queue
    return av_read_frame(input.format_ctx, packet);
}

int DecoderPacketData(AVCodecContext * decoder, AVPacket * packet, AVFrame * frame)
{
    int result = 0;
    while (result >= 0)
    {
        result = avcodec_receive_frame(decoder, frame);
        if (result == AVERROR(EAGAIN) || result == AVERROR_EOF)
        {
            result = 0;
            break;
        }
        else if (result < 0)
        {
            result = -1;
            break;
        }
        else
        {
            result = 1;
            break;
        }
    }
    return result;
}

int EncoderFrameData(AVCodecContext * encoder, AVFrame * frame, AVPacket * packet)
{
    int result = 0;
    while (result >= 0)
    {
        result = avcodec_receive_packet(encoder, packet);
        if (result == AVERROR(EAGAIN) || result == AVERROR_EOF)
        {
            result = 0;
            break;
        }
        else if (result < 0)
        {
            result = -1;
            break;
        }
        else
        {
            result = 1;
            break;
        }
    }
    return result;
}

int encode_write(InputFile * input, OutputFile * output, AVFrame *frame)
{
    int ret = 0;
    if ((ret = avcodec_send_frame(output->video_encoder_ctx, frame)) < 0)
    {
        fprintf(stderr, "Error code: %d.\n", ret);
        ret = -1;
    }
    while (ret >= 0) {
        ret = avcodec_receive_packet(output->video_encoder_ctx, output->m_encoder_packet);
        if (ret)
            break;
        DumpPacket(*input, *output, output->m_encoder_packet);
        output->m_encoder_packet->stream_index = 0;
        output->m_encoder_packet->dts = output->m_encoder_packet->pts;
        output->m_encoder_packet->duration = 0;

        av_interleaved_write_frame(output->format_ctx, output->m_encoder_packet);
        av_packet_unref(output->m_encoder_packet);
    }
    ret = ((ret == AVERROR(EAGAIN)) ? 0 : -1);
    return ret;
}

int EncoderAudioData(InputFile * input, OutputFile * output, AVFrame * frame)
{
    int result = avcodec_send_frame(output->audio_encoder_ctx, frame);
    if (result < 0)
    {
        return -1;
    }
    while (result >= 0)
    {
        result = avcodec_receive_packet(output->audio_encoder_ctx, output->m_encoder_packet);
        if (result)
        {
            break;
        }
        DumpPacket(*input, *output, output->m_encoder_packet);
        output->m_encoder_packet->stream_index = output->audio_encoder_index;
        output->m_encoder_packet->dts = output->m_encoder_packet->pts;
        output->m_encoder_packet->duration = 0;
        av_interleaved_write_frame(output->format_ctx, output->m_encoder_packet);
        av_packet_unref(output->m_encoder_packet);
    }
    result = ((result == AVERROR(EAGAIN)) ? 0 : -1);
    return result;
}

void ProcessVideo(InputFile & input, OutputFile & output, AVPacket * packet)
{
    if (packet) DumpPacket(input, output, packet);
    int result = avcodec_send_packet(input.video_decoder_ctx, packet);
    while (result >= 0)
    {
        result = avcodec_receive_frame(input.video_decoder_ctx, input.m_decoder_frame);
        if (result == AVERROR(EAGAIN) || result == AVERROR_EOF)
        {
            result = 0;
            break;
        }
        else if (result < 0)
        {
            result = -1;
            break;
        }
        //encode_write(&input, &output, input.m_decoder_frame);
        av_frame_unref(input.m_decoder_frame);
    }
    if (packet == nullptr && result >= 0)
    {
        //encode_write(&input, &output, NULL);
    }
}

void ProcessAudio(InputFile & input, OutputFile & output, AVPacket * packet)
{
    if (packet) DumpPacket(input, output, packet);
    int result = avcodec_send_packet(input.audio_decoder_ctx, packet);
    while (result >= 0)
    {
        result = avcodec_receive_frame(input.audio_decoder_ctx, input.m_decoder_frame);
        if (result == AVERROR(EAGAIN) || result == AVERROR_EOF)
        {
            result  = 0;
            break;
        }
        else if (result < 0)
        {
            result = -1;
            break;
        }
        //EncoderAudioData(&input, &output, input.m_decoder_frame);
        av_frame_unref(input.m_decoder_frame);
    }
    if (packet == NULL && result >= 0)
    {
        //EncoderAudioData(&input, &output, NULL);
    }
}

int main(int argc, const char * argv[])
{
    if (argc < 2)
    {
        printf("%s need input file.\n", __func__ );
        return EXIT_FAILURE;
    }
    const std::string input_filename = argv[1];
    std::string output_filename = argv[0];
    output_filename += ".mp4";

    VideoAudioParameter param_set;

    param_set.video_encoder_codec_name = "h.264";

    InputFile input;
    OutputFile output;

    input.filename = input_filename;
    output.output_filename = output_filename;

    if (OpenInputFile(input) < 0 || OpenOutputFile(input, output, param_set) < 0)
    {
        FreeInput(input);
        FreeOutput(output);
        return EXIT_FAILURE;
    }
    AVDictionary* opts = NULL;
    av_dict_set(&opts, "movflags", "frag_keyframe+empty_moov+default_base_moof", 0);
    avformat_write_header(output.format_ctx, &opts);
    do
    {
        if (ReadInputFrame(input, input.m_read_packet) < 0)
        {
            break;
        }
        AVStream * stream = input.format_ctx->streams[input.m_read_packet->stream_index];
        if (stream->codecpar->codec_type == AVMEDIA_TYPE_VIDEO)
        {
            //DumpPacket(output.m_dump_handle, packet);
            ProcessVideo(input, output, input.m_read_packet);
        }
        else
        {
            // process
            ProcessAudio(input, output, input.m_read_packet);
        }
        av_packet_unref(input.m_read_packet);

    } while (true);
    for (int i = 0; i < input.format_ctx->nb_streams; i++)
    {
        AVStream * stream = input.format_ctx->streams[i];
        if (stream->index == input.video_index)
        {
            ProcessVideo(input, output, NULL);
        }
        else if (stream->index == input.audio_index)
        {
            ProcessAudio(input, output, NULL);
        }
    }

    av_write_trailer(output.format_ctx);

    FreeInput(input);
    FreeOutput(output);
    return EXIT_SUCCESS;
}